Summary

Follow-up to #16 as suggested by @cgiovanetti. This PR replaces jnp.interp with interpax.interp1d in thermo.py for consistency with abundances.py and better performance.

Changes:

• Replace 6 jnp.interp calls with interpax.interp1d
• Flip QED correction tables at load time (instead of at runtime) since interpax requires monotonically increasing x coordinates
• Use constant extrapolation values (0.0, 1.0) for collision factor tables (matching original behavior)

Functions updated:

• rho_EM_std: QED corrections
• p_EM_std: QED correction
• rho_plus_p_EM_std: QED correction
• G_nue_with_me: Collision factor interpolation
• G_numt_with_me: Collision factor interpolation

Caveat

For the collision factor interpolations (G_nue_with_me, G_numt_with_me), the original code used left=f_tab[0,1], right=f_tab[-1,1] to extrapolate with boundary values. Since interpax.interp1d requires the extrap parameter to be concrete values (not traced JAX arrays), I hardcoded extrap=(0.0, 1.0) based on the current data files:

f_nue_ann  - first: 0.0, last: 1.0
f_nue_scat - first: 0.0, last: 1.0
f_numu_ann - first: 0.0, last: 1.0
f_numu_scat - first: 0.0, last: 1.0

If these data files are ever updated with different boundary values, the hardcoded extrapolation values would need to be updated accordingly.

Test plan

• pytest pytest/test_abundances.py passes
• Reference BBN abundances match expected values

🤖 Generated with Claude Code